package msu.ml.data;

import msu.ml.core.*;
import weka.core.*;
import java.io.*;
import java.util.*;

/**
 * The CrossValidationDataProvider divides the data or a subset
 * of the data in a directory into a number of folds for use in
 * a cross validation experiment. Using cross validation, if there
 * are n folds, then n-1 folds are used for training and the final 
 * fold is used for validation. This cycle repeats n times, thus all
 * files are classified exactly once.
 *
 * @author Reginald M Mead
 * @version 1.0
 */
public class CrossValidationDataProvider extends DataProvider 
{

    private String directoryPath;
    private String [][] folds;
    private int filesPerFold;
    private int numFolds;

    int currentFold;

    /**
     * Creates a new CrossValidationDataProvider
     *
     */
    public CrossValidationDataProvider()
    {
        filesPerFold = 0;
    }

    /**
     * Creates a new CrossValidationDataProvider using
     * the specified cache to load files
     *
     * @param cache DataCache used to load files
     */
    public CrossValidationDataProvider(DataCache cache)
    {
        super(cache);
        filesPerFold = 0;
    }

    public int count()
    {
        if(!this.isInitialized)
            initialize();

        return this.numFolds * this.filesPerFold;
    }


    /**
     * returns the directory path to the files.
     * @return directory path
     */
    public String getDirectoryPath()
    {
        return directoryPath;
    }

    /**
     * returns the number of files to be included in each fold.
     * @return number of files in each fold.
     */
    public int getFilesPerFold()
    {
        return filesPerFold;
    }

    /**
     * Provides the next group of target data instances.
     * Unlike the training data which can be aggregated into
     * one monolythic instances object, the target data
     * is provides as an array of instances, each representing
     * a single sweep. In this way the data can be classified
     * one sweep at a time.
     *
     * @return next group of target data instances
     */
    public NxInstances [] getNextTargetData()
    {
        if(!this.isInitialized)
            initialize();

        NxInstances [] result = new NxInstances[this.filesPerFold];
        for(int i = 0; i < this.filesPerFold; i++)
            result[i] = m_DataCache.retrieve(folds[this.currentFold][i]);

        this.currentFold++;

        return result;
    }

    /**
     * Provides the next set of training instances
     *
     * @return next training data
     */
    public NxInstances [] getNextIndividualTrainingData()
    {
        if(!this.isInitialized)
            initialize();

        NxInstances [] result = new NxInstances[this.filesPerFold * (this.numFolds - 1)];
        for(int i = 0, k = 0; i < this.numFolds; i++)
        {
            if(i != this.currentFold)
            {
                for(int j = 0; j < folds[i].length; j++, k++)
                {
                    result[k] = m_DataCache.retrieve(folds[i][j]);
                }
            }
        }

        return result;
    }

    /**
     * Provides the next set of training instances in one
     * large Instances object.
     *
     * @return next training data
     */
    public NxInstances getNextTrainingData()
    {
        if(!this.isInitialized)
            initialize();

        NxInstances result = null;
        for(int i = 0; i < this.numFolds; i++)
        {
            if(i != this.currentFold)
            {
                for(int j = 0; j < folds[i].length; j++)
                {
                    if(result == null)
                        result = m_DataCache.retrieve(folds[i][j]);
                    else
                        result.concatenate(m_DataCache.retrieve(folds[i][j]));
                }
            }
        }

        return result;
    }

    /**
     * Get the names of the data sources used in the
     * next set of training data.
     *
     * @return the labels of the next training data
     */
    public String [] getNextTrainingLabels()
    {
        if(!this.isInitialized)
            initialize();

        String [] result = new String[(this.numFolds - 1) * this.filesPerFold];
        for(int i = 0, k = 0; i < this.numFolds; i++)
            if(i != this.currentFold)
                for(int j = 0; j < folds[i].length; j++)
                    result[k++] = folds[i][j];

        return result;
    }

    /**
     * return the number of folds to divide the data into.
     *
     * @return number of folds
     */
    public int getNumFolds()
    {
        return numFolds;
    }

    /**
     * Return whether there is more data 
     * @return true if more data exists
     */
    public boolean hasMoreData()
    {
        if(!this.isInitialized)
            initialize();

        return this.currentFold < this.numFolds;
    }

    /**
     * Print a listing of the files in a given fold
     *
     * @param index index of fold to print
     */
    public void printFold(int index)
    {
        System.out.println();
        System.out.println("Fold #" + index);
        for(int i = 0; i < folds[index].length; i++)
        {
            System.out.println(folds[index][i]);
        }
    }

    /**
     * set the directory path to the files
     *
     * @param newVal directory path
     */
    public void setDirectoryPath(String newVal)
    {
        directoryPath = newVal;
    }

    /**
     * set the number of files per fold. Default is zero which
     * indicates that each fold should contain enough files to 
     * include the maximum number of files when all folds are combined.
     *
     * @param newVal number of files per fold
     */
    public void setFilesPerFold(int newVal)
    {
        filesPerFold = newVal;
    }

    /**
     * set the number of folds to use
     *
     * @param newVal number of folds to use
     */
    public void setNumFolds(int newVal)
    {
        numFolds = newVal;
    }

    protected boolean isInitialized = false;
    public void initialize()
    {
        File inputDirectory = new File(directoryPath);
        String [] files = inputDirectory.list();

        ArrayList<String> remainingFiles = new ArrayList<String>(files.length);
        //Add data files to list
        for(int i = 0; i < files.length; i++)
        {
            if(m_DataCache.canLoad(files[i]))
                remainingFiles.add(directoryPath + File.separatorChar + files[i]);
        }

        //Randomly remove excess files
        if(filesPerFold > 0)
        {
            while(remainingFiles.size() > filesPerFold * this.numFolds) 
            {
                int index = (int)(Math.random() * (double)remainingFiles.size());
                remainingFiles.remove(index);
            }
        }

        this.filesPerFold = remainingFiles.size() / this.numFolds;
        this.folds = new String[this.numFolds][this.filesPerFold];

        //Construct Folds
        for(int i = 0; i < this.numFolds; i++)
        {
            for(int j = 0; j < folds[i].length; j++)
            {
                int index = (int)(Math.random() * (double)remainingFiles.size());
                folds[i][j] = remainingFiles.remove(index);
            }
        }

        //Precache Data
        for(int i = 0; i < this.numFolds; i++)
        {
            for(int j = 0; j < folds[i].length; j++)
            {
                m_DataCache.preCache(folds[i][j]);
            }
        }

        this.currentFold = 0;

        this.isInitialized = true;
    }

}
